from pydantic import BaseModel,Field
from llmapi.base import *
from text2vec import SentenceModel
class Text2vecModel(BaseModel, EmbeddingModel):
    model_name: Optional[str] = Field(default="text2vec")
    model: SentenceModel
    class Config:
        arbitrary_types_allowed = True
    @property
    def name(self) -> str:
        return self.model_name or "text2vec"
    
    def embedding(self, input: str, **kwargs) -> Embedding:
        embeddings = self.model.encode(input)
        return {
            "object": "list",
            "model": self.name,
            "data": [
                {
                    "object": "embedding",
                    "embedding": embeddings.tolist(),
                    "index": 0
                }
            ],
            "usage": {
                "prompt_tokens": 0,
                "total_tokens": 0
            }
        }